{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66aabd37-663c-4f55-9635-cdf287c22232",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import datasets\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.transforms import ToTensor\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "from PIL import Image\n",
    "\n",
    "root = f\"{'/'.join(os.getcwd().split('/')[:3])}/hdfs/training-data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb3fc082-e59b-44b4-a9cc-99f45063c3ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_data = datasets.FashionMNIST(\n",
    "    root,\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root,\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "582cff69-6cd2-443e-bb77-851549dc8484",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "labels_map = {\n",
    "    0: \"T-Shirt\",\n",
    "    1: \"Trouser\",\n",
    "    2: \"Pullover\",\n",
    "    3: \"Dress\",\n",
    "    4: \"Coat\",\n",
    "    5: \"Sandal\",\n",
    "    6: \"Shirt\",\n",
    "    7: \"Sneaker\",\n",
    "    8: \"Bag\",\n",
    "    9: \"Ankle Boot\",\n",
    "}\n",
    "figure = plt.figure(figsize=(5, 5))  # figure 图片大小 5*5 inches(2.54cm)\n",
    "cols, rows = 3, 3\n",
    "for i in range(1, cols * rows + 1):\n",
    "    sample_idx = torch.randint(len(training_data), size=(1,)).item()\n",
    "    img, label = training_data[sample_idx]\n",
    "    figure.add_subplot(rows, cols, i)  # figure 中子图片的摆放在3*3的第i个位置\n",
    "    plt.title(labels_map[label])\n",
    "    plt.axis(\"off\")\n",
    "    plt.imshow(img.squeeze(), cmap=\"gray\")  # squeeze()挤压掉shapr=1的维度, 比如 [[[[1,2,1]]]] -> [1,2,1]\n",
    "    \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230b4e7e-2406-42ed-8a3a-08e9114f414c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# squeeze() 维度挤压\n",
    "arr = np.array([[[[2,1,4]], [[3,6,2]]]])\n",
    "brr = arr.squeeze()\n",
    "print(arr, arr.shape)\n",
    "print(brr, brr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96a1e2a3-bae6-45c8-b369-944e915d7ee9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cifar100 = datasets.CIFAR100(\n",
    "    root,\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e2a9c92-a778-4600-83dc-4c546641e985",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mytd = root + '/mytd'\n",
    "os.makedirs(mytd, exist_ok=True)\n",
    "index_file_path = os.path.join(mytd, 'index.txt') # 索引文件\n",
    "\n",
    "with open(index_file_path, 'w') as f:\n",
    "    \n",
    "    # 枚举张量, i为枚举值的索引, index为枚举值\n",
    "    size = len(cifar100)\n",
    "    i = 0\n",
    "    while i < 20:\n",
    "        img, label = cifar100[np.random.randint(size)]\n",
    "        if label < 10:\n",
    "            img = transforms.ToPILImage()(img)  # 将图片转换为 PIL 格式\n",
    "            filename = f'cifar-train-{i}.jpg'   # 生成文件名\n",
    "            img.save(os.path.join(mytd, filename))  # 保存图片\n",
    "            f.write(f'{filename}, {label}\\n') #  写索引\n",
    "            i += 1\n",
    "\n",
    "print(f'已保存 20 张图片到 {output_dir} 文件夹')\n",
    "print(f'索引文件已保存到 {index_file_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c167aa6-d633-4b69-83a2-9e1fa30a5bdc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 定义数据\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "from torchvision.io import read_image\n",
    "\n",
    "class CustomImageDataset(Dataset):\n",
    "    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):\n",
    "        self.img_labels = pd.read_csv(annotations_file)\n",
    "        print(len(self.img_labels))\n",
    "        self.img_dir = img_dir\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_labels)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])\n",
    "        image = read_image(img_path)  # 读取图片为 shape(3, n, n)数据 \n",
    "        label = self.img_labels.iloc[idx, 1]\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        if self.target_transform:\n",
    "            label = self.target_transform(label)\n",
    "        return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7878f2-ee4d-411f-9fd5-e6903147d6ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 加载并迭代自定义的训练数据\n",
    "def show_sample():\n",
    "    training_data = CustomImageDataset(index_file_path, mytd)\n",
    "    train_dataloader = DataLoader(training_data, batch_size=30, shuffle=True)\n",
    "    train_features, train_labels = next(iter(train_dataloader))\n",
    "    print(f\"Feature batch shape: {train_features.size()}\")\n",
    "    print(f\"Labels batch shape: {train_labels.size()}\")\n",
    "    img = train_features[0].squeeze()\n",
    "\n",
    "    # 在第三维堆叠张量, 因此输入张量维度最好小于3, 可用于图片数据转换3nn -> nn3\n",
    "    img = torch.dstack([img[0], img[1], img[2]])  \n",
    "    label = train_labels[0]\n",
    "    plt.figure(figsize=(0.4,0.4))\n",
    "    plt.imshow(img)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "    print(f\"Label: {label}\")\n",
    "\n",
    "show_sample()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312",
   "language": "python",
   "name": "py312"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
